pytorch自定义函数实现自动梯度

2023-09-12 23:46| 来源: 网络整理| 查看: 265

Motivation

构建模型有时需要使用自定义的函数，为了不影响模型的反向传播，需要实现自动梯度计算（即把自定义函数嵌入计算图）。

实现

要点：

将函数定义为类，需继承自torch.autograd.Function类需实现两个静态方法：forward()和backward()，分别对应前向传播和反向传播函数使用前需调用apply方法从而嵌入计算图，实现自动求导

用一个例子来说明: 假设我们要实现一个多项式拟合模型： y = a + b P 2 ( c x + d ) y = a + bP_2(cx + d) y=a+bP2(cx+d) ，用来拟合正切函数，其中 P 2 P_2 P2为一个二次函数： P 2 ( x ) = x 2 + 2 x P_2(x) = x^2 + 2x P2(x)=x2+2x，需要我们自定义实现

import torch import math dtype = torch.float #构建训练集,使用随机数据 x = torch.linspace(-math.pi, math.pi, 2000, dtype = dtype) y = torch.tan(x) #将权重值a,b,c,d随机初始化，设置requires_grad=True以实现自动求偏导 a = torch.randn((), dtype = dtype, requires_grad = True) b = torch.randn((), dtype = dtype, requires_grad = True) c = torch.randn((), dtype = dtype, requires_grad = True) d = torch.randn((), dtype = dtype, requires_grad = True) #实现P2 class P2(torch.autograd.Function): #定义forward方法，以在前向传播中实现函数的功能 @staticmethod def forward(ctx, input): #保存input值，反向传播计算梯度时会用到 ctx.save_for_backward(input) return input**2 + 2*input #定义backward方法 @staticmethod def backward(ctx,grad_output): #输入参数grad_output是目标函数对输出结果的梯度 #这里需要计算函数输出值相对于输入input的梯度u， #返回grad_output与u的乘积以实现链式法则 input, = ctx.saved_tensors return grad_output * (2 * inptu + 2) learning_rate = 1e-6 #全量学习 for t in range(2000): #对刚刚定义的函数fun调用apply方法，使其嵌入计算图 p2 = P2.apply #前向传播 y_pred = a + b * p2(c * x + d) #定义损失函数 loss = (y_pred - y).pow(2).sum() if t%100 == 99: print(t, loss.item()) #反向传播计算梯度 loss.backward() #更新参数 #这里每个参数的梯度自定在反向传播时自动计算好并保存在param.grad中了 with torch.no_grad(): a -= learning_rate * a.grad b -= learning_rate * b.grad c -= learning_rate * c.grad d -= learning_rate * d.grad #将梯度重置为0 a.grad = b.grad = c.grad = d.grad = None print(f'Result: y = {a.item()} + {b.item()} * P2({c.item} x + {d.item()}')

运行结果：在这里插入图片描述（由于参数选择很随意，且权重是随机初始化的，所以拟合效果不好）

参考资料

链接：Learning pytorch with examples

【本文地址】

公司简介

联系我们